// all trax programs must include trax.hpp
#include "trax.hpp"
#include "Image.h"
#include "Vector3.h"
#include "Sphere.h"
#include "HitRecord.h"
#include "PinHoleCamera.h"
#include "Scene.h"
#include "Shader.h"
#include "AABBox.h"
#include "Unit_Test.h"
#include "atTriangle.h"
#include "RayTracer.h"
#include "atUtil.h"
#include "atMaterial.h"
#include "PhotonMap.h"
#include "atAtomicReg.h"
#include "HitMap.h"

#define MAX_MSAA 128
#define PHT_RADIUS 1.0f
#define ALPHA_R 0.5f
#define PHT_PASS 1

#define ONLY_PHOTON_MAP 0
#define RUN_TEST 0


// "main" function is called "trax_main", returns void instead of int
void trax_main()
{
    using namespace AtomTrace;
#if RUN_TEST
    printf("Unit test\n");
	//Unit_test_Vector3();
    //Unit_test_BoxIntersection();
    //Unit_test_TriangleIntersection();
    //Unit_test_BVH_Stack();
    //Unit_test_GetMemInfo();
    //Unit_test_Atomicadd();
    Unit_test_PosToZcurve();
    return;
#endif
    
    const int threadId = atomicinc(REG_THREAD_ID);

    const int maxDepth = 7;

    float color_threshold = 0.001f;
	int xres = GetXRes();
	int yres = GetYRes();
	int start_fb = GetFrameBuffer();
    int num_samples = loadi(TRAX_NUM_SAMPLES);
    float num_samplesF_inv = 1.0f / static_cast<float>(num_samples);
	Image image(xres, yres, start_fb);
    HitMap hitmap;
	

    //<Scene def. lab3>------------------
    PinHoleCamera pcamera;
    PointLight light;
    //Vector3 lpos;
    Scene scene;
    atTriangle tempTri;

    scene.numLights = 1;
    //light..LoadFromGM(GetLight());
    //scene.pointLights = &light;
    scene.startLightAddr = GetLight();
    //lpos.LoadFromGM(GetLight());
    scene.startMatAddr = GetMaterials();
    pcamera.LoadFromGM(GetCamera());
    scene.startTriAddr = GetStartTriangles();
    scene.numTri = GetNumTriangles();
    scene.bvh_start = GetBVH();
    //-----------------------</Scene def. lab3>
    
    //<PhotonMap test> ------------------------------
    PhotonMap pm;
    pm.SetMaterialStartAddr(scene.startMatAddr);
    pm.Compute(scene, threadId);

#if ONLY_PHOTON_MAP
    pm.Reset();
    	if(threadId==0)
    		printf("--------------- No rendering!! ----------------\n");
        return;
#endif
    //--------------------------------</PhotonMap test>*/

    //light.LoadFromGM(scene.startLightAddr);
    //printf("lpos: [%f, %f, %f]\n", light.pos[0], light.pos[1], light.pos[2]);

    int imgSize = xres * yres;
    int subImg = 1;
    int bImgSize = imgSize/subImg;
    //Vector3 des;
    //int nodeID = 0;
    int percent = 0;
    Color res, red, sumres;
    Color tresult[MAX_MSAA];
    //red.Set(1.0f, 0.0f, 0.0f);

    // ---------- Loop through photon pass --------------------
    for (int ppass=0; ppass < PHT_PASS; ppass++) {
        printf("thread %d, photon pass %d\n", threadId, ppass);
    
        // ---------- Loop through every pixels in z-shape pattern ---------
        for(int pixww = atomicinc(REG_PIXEL_ID)*subImg;
                pixww < imgSize;
                pixww = atomicinc(REG_PIXEL_ID) * subImg)
        //int pixww = 63; //256+512 * 256;
        {
            float pecf = static_cast<float>(pixww) * 100.0f / static_cast<float>(imgSize);
            printf("\rrendering %.2f", pecf);
            for(int subPix = 0; subPix<subImg; subPix++)
            {
                int pix = pixww+subPix;
                int ix, iy;
                ConvertLineIndexToZindexXY(pix, ix, iy);
                
                float offx = 0.0f;
                float offy = 0.0f;
                Ray ray;
                int outpix;
                outpix = iy * xres + ix;
                res.Set(0.0f, 0.0f, 0.0f);

                // ---- multi sampling anti-aliasing ---------
                int num_used_samples = num_samples;
                Color average;
                // TODO: progressive photon map
                float maxR2;
                HitPoint hp_ave;
                hitmap.Get(outpix, hp_ave);
                
                
                HitPoint curHit;
                HitPoint newHp;
                if(ppass == 0)
                {
                    hp_ave.r = PHT_RADIUS; // Initial estimate radius
                    curHit.r = PHT_RADIUS; // Initial estimate radius
                    newHp.r = PHT_RADIUS;
                }
                
                hp_ave.pix = outpix;
                curHit.pix = outpix;
                newHp.pix = outpix;
                //for(int i=0; i<num_samples; i++)
                {
                    int i=0;
                    //offx = (trax_rand() - 0.5f);
                    //offy = (trax_rand() - 0.5f);
                    
                    // Halton sequence offset
                    offx = Halton(i, 2);
                    offy = Halton(i, 3);

                    // Center them
                    offx += (offx>0.5f)? -1.0f : 0.0f;
                    offy += (offy>0.5f)? -1.0f : 0.0f;

                    ray = pcamera.CastRay(ix+offx, iy+offy);
                    int depth = 0;

                    HitRecord hit;
                    
                    tresult[i].Set(0.0f, 0.0f, 0.0f);
                    while(depth == 0 || (depth < maxDepth && hit.newDir.LengthSquare() > 0.1f))
                    {
                        depth++;
                        hit.Init();
                        bool isHit = false;
                        isHit = RayTracer::TraceBVH(ray, hit);
                        if(isHit)
                        {
                            //tresult[i].Set(1.0f, 1.0f, 1.0f);
                            //tresult[i] += shade(ray, hit, scene);
                            
                            Color pht_flux;
                            pm.RadianceEstimate(hit, curHit.r, pht_flux, maxR2, newHp.npht);
                            
                            pht_flux *= 1.0f / (3.141592653f * maxR2);
                            tresult[i] += pht_flux;
                            
                            int illum = hit.matIllum;
                            Material::GetNewDir(illum, hit, hit.newDir, false);

                            ray._o = hit.P + hit.newDir.GetNormalize() * 0.01f;
                            ray._dir = hit.newDir;
                        }else
                        {
                            tresult[i] += scene.background;
                            break;
                        }
                        //tresult[i] *= 0.8f;
                    } // bounce rays
                    
                    tresult[i].Clamp(0.0f, 1.0f);
                    res += tresult[i];

                    // Adaptive multi-sampling
                    average = res * (1.0f / static_cast<float>(i+1));
                    Color variance;
                    variance.Set(0.0f, 0.0f, 0.0f);
                    if(i>=4){
                        for(int j=0; j<i; j++)
                        {
                            Color tempDiff = tresult[j] - average;
                            variance += tempDiff * tempDiff;
                        }
                        variance *= (1.0f / static_cast<float>(i+1));
                        if((variance.r() < color_threshold && 
                            variance.g() < color_threshold &&
                            variance.b() < color_threshold) || 
                            i == num_samples-1)
                        {
                            //res = average;
                            num_used_samples = i+1;

                            break;

                        }
                    } // anti aliasing estimation
                } // loop through number of samples, anti-aliasing

                //res *= 1.0f / static_cast<float>(num_used_samples);
                image.set(outpix, average);
            } // loop through pixels in block
        } // loop through pixels block
        
        barrier(REG_BARRIER);
        if(threadId == 0)
        {
            int pixId = atomicadd(REG_PIXEL_ID, 0);
            pixId = atomicadd(REG_PIXEL_ID, -pixId);
            printf("Rest register for pixel_ID: %d\n", pixId);
        }
        barrier(REG_BARRIER);
    } // loop through photon pass
    
    barrier(REG_BARRIER);
    if(threadId == 0)
    {
        int pixId = atomicadd(REG_PIXEL_ID, 0);
        pixId = atomicadd(REG_PIXEL_ID, -pixId);
        printf("Rest register for pixel_ID: %d\n", pixId);
    }
    barrier(REG_BARRIER);
    
    /*/ Copy hit map result to frame buffer
    for (int pix = atomicinc(REG_PIXEL_ID); pix<imgSize; pix = atomicinc(REG_PIXEL_ID))
    {
        HitPoint hit;
        hitmap.Get(pix, hit);
        image.set(pix, hit.flux);
    } // Copy hit map result to frame buffer */
}


